/// @file
/// Calculation of biinvariant means.

#ifndef SOPHUS_AVERAGE_HPP
#define SOPHUS_AVERAGE_HPP

#include "common.hpp"
#include "rxso2.hpp"
#include "rxso3.hpp"
#include "se2.hpp"
#include "se3.hpp"
#include "sim2.hpp"
#include "sim3.hpp"
#include "so2.hpp"
#include "so3.hpp"

namespace Sophus
{
  // Calculates mean iteratively.
  //
  // Returns ``nullopt`` if it does not converge.
  //
  template <class SequenceContainer>
  optional<typename SequenceContainer::value_type> iterativeMean(
      SequenceContainer const &foo_Ts_bar, int max_num_iterations)
  {
    size_t N = foo_Ts_bar.size();
    SOPHUS_ENSURE(N >= 1, "N must be >= 1.");

    using Group = typename SequenceContainer::value_type;
    using Scalar = typename Group::Scalar;
    using Tangent = typename Group::Tangent;

    // This implements the algorithm in the beginning of Sec. 4.2 in
    // ftp://ftp-sop.inria.fr/epidaure/Publications/Arsigny/arsigny_rr_biinvariant_average.pdf.
    Group foo_T_average = foo_Ts_bar.front();
    Scalar w = Scalar(1. / N);
    for (int i = 0; i < max_num_iterations; ++i)
    {
      Tangent average;
      setToZero<Tangent>(average);
      for (Group const &foo_T_bar : foo_Ts_bar)
      {
        average += w * (foo_T_average.inverse() * foo_T_bar).log();
      }

      Group foo_T_newaverage = foo_T_average * Group::exp(average);
      if (squaredNorm<Tangent>(
              (foo_T_newaverage.inverse() * foo_T_average).log()) <
          Constants<Scalar>::epsilon())
      {
        return foo_T_newaverage;
      }

      foo_T_average = foo_T_newaverage;
    }

    // LCOV_EXCL_START
    return nullopt;
    // LCOV_EXCL_STOP
  }

#ifdef DOXYGEN_SHOULD_SKIP_THIS
  // Mean implementation for any Lie group.
  template <class SequenceContainer, class Scalar>
  optional<typename SequenceContainer::value_type> average(
      SequenceContainer const &foo_Ts_bar);
#else

  // Mean implementation for SO(2).
  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, SO2<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar)
  {
    // This implements rotational part of Proposition 12 from Sec. 6.2 of
    // ftp://ftp-sop.inria.fr/epidaure/Publications/Arsigny/arsigny_rr_biinvariant_average.pdf.
    size_t N = std::distance(std::begin(foo_Ts_bar), std::end(foo_Ts_bar));
    SOPHUS_ENSURE(N >= 1, "N must be >= 1.");
    SO2<Scalar> foo_T_average = foo_Ts_bar.front();
    Scalar w = Scalar(1. / N);

    Scalar average(0);
    for (SO2<Scalar> const &foo_T_bar : foo_Ts_bar)
    {
      average += w * (foo_T_average.inverse() * foo_T_bar).log();
    }

    return foo_T_average * SO2<Scalar>::exp(average);
  }

  // Mean implementation for RxSO(2).
  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, RxSO2<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar)
  {
    size_t N = std::distance(std::begin(foo_Ts_bar), std::end(foo_Ts_bar));
    SOPHUS_ENSURE(N >= 1, "N must be >= 1.");
    RxSO2<Scalar> foo_T_average = foo_Ts_bar.front();
    Scalar w = Scalar(1. / N);

    Vector2<Scalar> average(Scalar(0), Scalar(0));
    for (RxSO2<Scalar> const &foo_T_bar : foo_Ts_bar)
    {
      average += w * (foo_T_average.inverse() * foo_T_bar).log();
    }

    return foo_T_average * RxSO2<Scalar>::exp(average);
  }

  namespace details
  {
    template <class T>
    void getQuaternion(T const &);

    template <class Scalar>
    Eigen::Quaternion<Scalar> getUnitQuaternion(SO3<Scalar> const &R)
    {
      return R.unit_quaternion();
    }

    template <class Scalar>
    Eigen::Quaternion<Scalar> getUnitQuaternion(RxSO3<Scalar> const &sR)
    {
      return sR.so3().unit_quaternion();
    }

    template <class SequenceContainer,
              class Scalar = typename SequenceContainer::value_type::Scalar>
    Eigen::Quaternion<Scalar> averageUnitQuaternion(
        SequenceContainer const &foo_Ts_bar)
    {
      // This:  http://stackoverflow.com/a/27410865/1221742
      size_t N = std::distance(std::begin(foo_Ts_bar), std::end(foo_Ts_bar));
      SOPHUS_ENSURE(N >= 1, "N must be >= 1.");
      Eigen::Matrix<Scalar, 4, Eigen::Dynamic> Q(4, N);
      int i = 0;
      Scalar w = Scalar(1. / N);
      for (auto const &foo_T_bar : foo_Ts_bar)
      {
        Q.col(i) = w * details::getUnitQuaternion(foo_T_bar).coeffs();
        ++i;
      }

      Eigen::Matrix<Scalar, 4, 4> QQt = Q * Q.transpose();
      // TODO: Figure out why we can't use SelfAdjointEigenSolver here.
      Eigen::EigenSolver<Eigen::Matrix<Scalar, 4, 4>> es(QQt);

      std::complex<Scalar> max_eigenvalue = es.eigenvalues()[0];
      Eigen::Matrix<std::complex<Scalar>, 4, 1> max_eigenvector =
          es.eigenvectors().col(0);

      for (int i = 1; i < 4; i++)
      {
        if (std::norm(es.eigenvalues()[i]) > std::norm(max_eigenvalue))
        {
          max_eigenvalue = es.eigenvalues()[i];
          max_eigenvector = es.eigenvectors().col(i);
        }
      }

      Eigen::Quaternion<Scalar> quat;
      quat.coeffs() << max_eigenvector[0].real(),
          max_eigenvector[1].real(),
          max_eigenvector[2].real(),
          max_eigenvector[3].real();

      return quat;
    }
  } // namespace details

  // Mean implementation for SO(3).
  //
  // TODO: Detect degenerated cases and return nullopt.
  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, SO3<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar)
  {
    return SO3<Scalar>(details::averageUnitQuaternion(foo_Ts_bar));
  }

  // Mean implementation for R x SO(3).
  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, RxSO3<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar)
  {
    size_t N = std::distance(std::begin(foo_Ts_bar), std::end(foo_Ts_bar));

    SOPHUS_ENSURE(N >= 1, "N must be >= 1.");
    Scalar scale_sum = Scalar(0);
    using std::exp;
    using std::log;
    for (RxSO3<Scalar> const &foo_T_bar : foo_Ts_bar)
    {
      scale_sum += log(foo_T_bar.scale());
    }

    return RxSO3<Scalar>(exp(scale_sum / Scalar(N)),
                         SO3<Scalar>(details::averageUnitQuaternion(foo_Ts_bar)));
  }

  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, SE2<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar, int max_num_iterations = 20)
  {
    // TODO: Implement Proposition 12 from Sec. 6.2 of
    // ftp://ftp-sop.inria.fr/epidaure/Publications/Arsigny/arsigny_rr_biinvariant_average.pdf.
    return iterativeMean(foo_Ts_bar, max_num_iterations);
  }

  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, Sim2<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar, int max_num_iterations = 20)
  {
    return iterativeMean(foo_Ts_bar, max_num_iterations);
  }

  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, SE3<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar, int max_num_iterations = 20)
  {
    return iterativeMean(foo_Ts_bar, max_num_iterations);
  }

  template <class SequenceContainer,
            class Scalar = typename SequenceContainer::value_type::Scalar>
  enable_if_t<
      std::is_same<typename SequenceContainer::value_type, Sim3<Scalar>>::value,
      optional<typename SequenceContainer::value_type>>
  average(SequenceContainer const &foo_Ts_bar, int max_num_iterations = 20)
  {
    return iterativeMean(foo_Ts_bar, max_num_iterations);
  }

#endif // DOXYGEN_SHOULD_SKIP_THIS

} // namespace Sophus

#endif // SOPHUS_AVERAGE_HPP
